from web.ais.TengXunYuanBao import JiChuBan
import requests
import hashlib
import hmac
import json
import os
from urllib.parse import quote
import xml.etree.ElementTree as ET


# 上传图片类
class TengXunYuanBao_ImageUpdate(JiChuBan.YuanBaoJiChuBan):
    def __init__(self):
        super().__init__()

    
    # 第一个请求(获取信息)
    def get_user_agent(self, img_path):
        json_data = {
            'fileName': img_path,
            'docFrom': 'localDoc',
        }

        response = requests.post(f'{self.domain}/api/resource/genUploadInfo', cookies=self.cookies, headers=self.headers, json=json_data)
        return response.json()


    # 第二个请求(上传图片)
    def upload_image(self, img_path):
        a1_json = self.get_user_agent(img_path)
        encryptTmpSecretId = a1_json['encryptTmpSecretId']
        location = a1_json['location']
        encryptToken = a1_json['encryptToken']
        startTime = a1_json['startTime']
        expiredTime = a1_json['expiredTime']
        singtrue = self.get_singtrue(a1_json, img_path)
        headers = {
        'Accept': '*/*',
        'Accept-Language': 'zh-CN,zh;q=0.9,en;q=0.8,en-GB;q=0.7,en-US;q=0.6',
        'Authorization': f'q-sign-algorithm=sha1&q-ak={encryptTmpSecretId}&q-sign-time={startTime};{expiredTime}&q-key-time={startTime};{expiredTime}&q-header-list=content-length;host;pic-operations&q-url-param-list=&q-signature={singtrue}',
        'Cache-Control': '',
        'Connection': 'keep-alive',
        'Content-Type': '',
        'Origin': self.domain,
        'Pic-Operations': '{"is_pic_info":1,"rules":[{"fileid":"' + location +'","rule":"imageMogr2/format/jpg"}]}',
        'Pragma': 'no-cache',
        'Referer': self.domain,
        'Sec-Fetch-Dest': 'empty',
        'Sec-Fetch-Mode': 'cors',
        'Sec-Fetch-Site': 'cross-site',
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36 Edg/133.0.0.0',
        'sec-ch-ua': '"Not(A:Brand";v="99", "Microsoft Edge";v="133", "Chromium";v="133"',
        'sec-ch-ua-mobile': '?0',
        'sec-ch-ua-platform': '"Windows"',
        'x-cos-security-token': encryptToken,
        }
        data = open(img_path, 'rb').read()

        response = requests.put(f'https://{a1_json['bucketName']}.{a1_json['accelerateDomain']}/{a1_json['location']}', headers=headers, data=data)
        return {
            "xml": self.xml_to_dict(ET.fromstring(response.text)),
            "response": a1_json
        }
    
    def get_singtrue(self, genUploadInfo, img_path):
        times = str(genUploadInfo['startTime'])+';'+ str(genUploadInfo['expiredTime'])
        b = self.hmac_sha1(times,genUploadInfo['encryptTmpSecretKey'])
        x = '\n'.join(['put', genUploadInfo['location'], '', self.obj2str(genUploadInfo, img_path), ''])
        s = '\n'.join(["sha1", times, self.sha1_string(x), ""])
        singtrue = self.hmac_sha1(s,b)
        return singtrue

    # 获取图片大小
    def get_img_size(self,img_path):

        # 图片文件的路径
        image_path = img_path

        try:
            # 获取文件大小（字节）
            file_size = os.path.getsize(image_path)
            return file_size
        except FileNotFoundError:
            print(f"文件 {image_path} 未找到。")



    def hmac_sha1(self,value, key):
        # 确保 value 和 key 是字节类型
        if isinstance(value, str):
            value = value.encode('utf-8')
        if isinstance(key, str):
            key = key.encode('utf-8')
        # 创建 HMAC 对象，使用 SHA1 算法
        h = hmac.new(key, value, hashlib.sha1)
        # 获取十六进制的哈希值
        hash_hex = h.hexdigest()
        return hash_hex



    def sha1_string(self,x):
        # 创建 SHA1 哈希对象
        sha1 = hashlib.sha1()

        # 由于 update 方法需要字节类型的数据，因此对字符串进行编码
        sha1.update(x.encode('utf-8'))
        return sha1.hexdigest()


    def obj2str(self,genUploadInfo, img_path):
        ContentLength = self.get_img_size(img_path)
        host = genUploadInfo['bucketName'] + "." + genUploadInfo['accelerateDomain']
        PicOperations = json.dumps({
            "is_pic_info": 1,
            "rules": [
                {
                    "fileid": genUploadInfo['location'],
                    "rule": "imageMogr2/format/jpg"
                }
            ]
        })
        # 去除PicOperations的空格
        PicOperations = PicOperations.replace(" ", "")
        result = f"content-length={ContentLength}&host={host}&pic-operations={quote(PicOperations,encoding='utf-8', safe='')}"
        return result


    # 开始
    def start(self, img_path):
        # 判断是本地图片是否存在
        if os.path.exists(img_path):
            return {
                "code": 200,
                "message": "图片存在",
                "data": self.upload_image(img_path)
            }
        else:
            return {
                "code": 404,
                "message": "图片不存在"
            }

    def xml_to_dict(self,element):
        result = {}
        # 遍历当前元素的子元素
        for child in element:
            if len(child) == 0:
                # 如果子元素没有子元素，直接将其标签和文本作为键值对添加到结果字典中
                result[child.tag] = child.text
            else:
                # 如果子元素有子元素，递归调用 xml_to_dict 函数处理子元素
                result[child.tag] = self.xml_to_dict(child)
        return result


# 图片编辑类
class TengXunYuanBao_ImageEdit(TengXunYuanBao_ImageUpdate):
    def __init__(self):
        super().__init__()
    

    # 风格转换
    def style_conversion(self, img_path, style):
        data = {"imageUrl":img_path,"prompt":f"转换为{style}风格","style":style + "风格","agentId":""}
        data = json.dumps(data)
        response = requests.post(f'{self.domain}/api/image/style', cookies=self.cookies, headers=self.headers, data=data, stream=True)
        self.handle_stream_response(response)
        return response
    
    # 去除水印
    def remove_watermark(self, img_path):
        data = '{"imageUrl":"' + img_path + '","agentId":""}'
        response = requests.post(f'{self.domain}/api/image/removewatermark', cookies=self.cookies, headers=self.headers, data=data, stream=True)
        self.handle_stream_response(response)
        return response


    # 变清晰
    def sharpen(self, img_path):
        data = '{"imageUrl":"' + img_path + '","agentId":""}'
        response = requests.post(f'{self.domain}/api/image/clarity', cookies=self.cookies, headers=self.headers, data=data, stream=True)
        self.handle_stream_response(response)
        return response

if __name__ == "__main__":
    image_edit = TengXunYuanBao_ImageEdit()





